{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "SocraticModels-MSR-VTT.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "8T7vyGBaRxPS",
        "sdGpKAVDOuAg",
        "cGZfnNoTOyI1"
      ],
      "machine_shape": "hm"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eE_aQ8yOVtkh"
      },
      "source": [
        "Copyright 2021 Google LLC.\n",
        "SPDX-License-Identifier: Apache-2.0\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "you may not use this file except in compliance with the License.\n",
        "You may obtain a copy of the License at\n",
        "\n",
        "https://www.apache.org/licenses/LICENSE-2.0\n",
        "\n",
        "Unless required by applicable law or agreed to in writing, software\n",
        "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "See the License for the specific language governing permissions and\n",
        "limitations under the License.\n",
        "\n",
        "# **Socratic Models: MSR-VTT Video-to-Text Retrieval**\n",
        "\n",
        "Socratic Models (SMs) is a framework that composes multiple pre-existing foundation models (e.g., large language models, visual language models, audio-language models) to provide results for new multimodal tasks, without any model finetuning.\n",
        "\n",
        "This colab runs SMs for zero-shot video-to-text retrieval on the [MSR-VTT](https://www.microsoft.com/en-us/research/wp-content/uploads/2016/06/cvpr16.msr-vtt.tmei_-1.pdf) Full and 1k-A test sets. Specifically, this augments [Portillo-Quintero et al. 2021](https://arxiv.org/pdf/2102.12443.pdf) with audio information by using an ALM for speech-to-text, summarizing the transcriptions with a causal LM (e.g., GPT-3), and re-ranking CLIP (VLM) matching scores against captions with a masked LM (e.g., RoBERTa) on the summaries.\n",
        "\n",
        "This is a reference implementation of one task demonstrated in the work: [Socratic Models: Composing Zero-Shot Multimodal Reasoning with Language](https://socraticmodels.github.io/)\n",
        "\n",
        "**Disclaimer:** this colab uses CLIP and GPT-3 as foundation models, and may be subject to unwanted biases. This code should be used with caution (and checked for correctness) in downstream applications.\n",
        "\n",
        "### **Quick Start:**\n",
        "\n",
        "**Step 1.** Register for an [OpenAI API key](https://openai.com/blog/openai-api/) to use GPT-3 (there's a free trial) and enter it below\n",
        "\n",
        "**Step 2.** Menu > Change runtime type > Hardware accelerator > \"GPU\"\n",
        "\n",
        "**Step 3.** Menu > Runtime > Run all"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "openai_api_key = \"your-api-key\""
      ],
      "metadata": {
        "id": "M4TTzIk0RsTm"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## **Setup**\n",
        "This installs a few dependencies: PyTorch, CLIP, GPT-3."
      ],
      "metadata": {
        "id": "8T7vyGBaRxPS"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pLBWqlYjXZsN",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "e5d2a87a-5661-4ba6-8484-5241261d2304"
      },
      "source": [
        "!pip install -U --no-cache-dir gdown --pre\n",
        "!pip install -U sentence-transformers\n",
        "!pip install openai ftfy\n",
        "!nvidia-smi  # Show GPU info."
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: gdown in /usr/local/lib/python3.7/dist-packages (4.4.0)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.7/dist-packages (from gdown) (4.62.3)\n",
            "Requirement already satisfied: requests[socks] in /usr/local/lib/python3.7/dist-packages (from gdown) (2.23.0)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.7/dist-packages (from gdown) (3.4.2)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.7/dist-packages (from gdown) (1.15.0)\n",
            "Requirement already satisfied: beautifulsoup4 in /usr/local/lib/python3.7/dist-packages (from gdown) (4.6.3)\n",
            "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from requests[socks]->gdown) (3.0.4)\n",
            "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.7/dist-packages (from requests[socks]->gdown) (1.24.3)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests[socks]->gdown) (2021.10.8)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests[socks]->gdown) (2.10)\n",
            "Requirement already satisfied: PySocks!=1.5.7,>=1.5.6 in /usr/local/lib/python3.7/dist-packages (from requests[socks]->gdown) (1.7.1)\n",
            "Requirement already satisfied: sentence-transformers in /usr/local/lib/python3.7/dist-packages (2.2.0)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from sentence-transformers) (1.19.5)\n",
            "Requirement already satisfied: torch>=1.6.0 in /usr/local/lib/python3.7/dist-packages (from sentence-transformers) (1.10.0+cu111)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.7/dist-packages (from sentence-transformers) (4.62.3)\n",
            "Requirement already satisfied: scipy in /usr/local/lib/python3.7/dist-packages (from sentence-transformers) (1.4.1)\n",
            "Requirement already satisfied: sentencepiece in /usr/local/lib/python3.7/dist-packages (from sentence-transformers) (0.1.96)\n",
            "Requirement already satisfied: nltk in /usr/local/lib/python3.7/dist-packages (from sentence-transformers) (3.2.5)\n",
            "Requirement already satisfied: torchvision in /usr/local/lib/python3.7/dist-packages (from sentence-transformers) (0.11.1+cu111)\n",
            "Requirement already satisfied: scikit-learn in /usr/local/lib/python3.7/dist-packages (from sentence-transformers) (1.0.2)\n",
            "Requirement already satisfied: transformers<5.0.0,>=4.6.0 in /usr/local/lib/python3.7/dist-packages (from sentence-transformers) (4.17.0)\n",
            "Requirement already satisfied: huggingface-hub in /usr/local/lib/python3.7/dist-packages (from sentence-transformers) (0.4.0)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from torch>=1.6.0->sentence-transformers) (3.10.0.2)\n",
            "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.7/dist-packages (from transformers<5.0.0,>=4.6.0->sentence-transformers) (21.3)\n",
            "Requirement already satisfied: tokenizers!=0.11.3,>=0.11.1 in /usr/local/lib/python3.7/dist-packages (from transformers<5.0.0,>=4.6.0->sentence-transformers) (0.11.6)\n",
            "Requirement already satisfied: importlib-metadata in /usr/local/lib/python3.7/dist-packages (from transformers<5.0.0,>=4.6.0->sentence-transformers) (4.10.0)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.7/dist-packages (from transformers<5.0.0,>=4.6.0->sentence-transformers) (3.4.2)\n",
            "Requirement already satisfied: sacremoses in /usr/local/lib/python3.7/dist-packages (from transformers<5.0.0,>=4.6.0->sentence-transformers) (0.0.49)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.7/dist-packages (from transformers<5.0.0,>=4.6.0->sentence-transformers) (2.23.0)\n",
            "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.7/dist-packages (from transformers<5.0.0,>=4.6.0->sentence-transformers) (2019.12.20)\n",
            "Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.7/dist-packages (from transformers<5.0.0,>=4.6.0->sentence-transformers) (6.0)\n",
            "Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /usr/local/lib/python3.7/dist-packages (from packaging>=20.0->transformers<5.0.0,>=4.6.0->sentence-transformers) (3.0.6)\n",
            "Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.7/dist-packages (from importlib-metadata->transformers<5.0.0,>=4.6.0->sentence-transformers) (3.7.0)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.7/dist-packages (from nltk->sentence-transformers) (1.15.0)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests->transformers<5.0.0,>=4.6.0->sentence-transformers) (2.10)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests->transformers<5.0.0,>=4.6.0->sentence-transformers) (2021.10.8)\n",
            "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from requests->transformers<5.0.0,>=4.6.0->sentence-transformers) (3.0.4)\n",
            "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.7/dist-packages (from requests->transformers<5.0.0,>=4.6.0->sentence-transformers) (1.24.3)\n",
            "Requirement already satisfied: click in /usr/local/lib/python3.7/dist-packages (from sacremoses->transformers<5.0.0,>=4.6.0->sentence-transformers) (7.1.2)\n",
            "Requirement already satisfied: joblib in /usr/local/lib/python3.7/dist-packages (from sacremoses->transformers<5.0.0,>=4.6.0->sentence-transformers) (1.1.0)\n",
            "Requirement already satisfied: threadpoolctl>=2.0.0 in /usr/local/lib/python3.7/dist-packages (from scikit-learn->sentence-transformers) (3.0.0)\n",
            "Requirement already satisfied: pillow!=8.3.0,>=5.3.0 in /usr/local/lib/python3.7/dist-packages (from torchvision->sentence-transformers) (7.1.2)\n",
            "Requirement already satisfied: openai in /usr/local/lib/python3.7/dist-packages (0.15.0)\n",
            "Requirement already satisfied: ftfy in /usr/local/lib/python3.7/dist-packages (6.1.1)\n",
            "Requirement already satisfied: requests>=2.20 in /usr/local/lib/python3.7/dist-packages (from openai) (2.23.0)\n",
            "Requirement already satisfied: pandas-stubs>=1.1.0.11 in /usr/local/lib/python3.7/dist-packages (from openai) (1.2.0.49)\n",
            "Requirement already satisfied: openpyxl>=3.0.7 in /usr/local/lib/python3.7/dist-packages (from openai) (3.0.9)\n",
            "Requirement already satisfied: pandas>=1.2.3 in /usr/local/lib/python3.7/dist-packages (from openai) (1.3.5)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.7/dist-packages (from openai) (4.62.3)\n",
            "Requirement already satisfied: et-xmlfile in /usr/local/lib/python3.7/dist-packages (from openpyxl>=3.0.7->openai) (1.1.0)\n",
            "Requirement already satisfied: numpy>=1.17.3 in /usr/local/lib/python3.7/dist-packages (from pandas>=1.2.3->openai) (1.19.5)\n",
            "Requirement already satisfied: python-dateutil>=2.7.3 in /usr/local/lib/python3.7/dist-packages (from pandas>=1.2.3->openai) (2.8.2)\n",
            "Requirement already satisfied: pytz>=2017.3 in /usr/local/lib/python3.7/dist-packages (from pandas>=1.2.3->openai) (2018.9)\n",
            "Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.7/dist-packages (from pandas-stubs>=1.1.0.11->openai) (3.10.0.2)\n",
            "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.7/dist-packages (from python-dateutil>=2.7.3->pandas>=1.2.3->openai) (1.15.0)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests>=2.20->openai) (2021.10.8)\n",
            "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from requests>=2.20->openai) (3.0.4)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests>=2.20->openai) (2.10)\n",
            "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.7/dist-packages (from requests>=2.20->openai) (1.24.3)\n",
            "Requirement already satisfied: wcwidth>=0.2.5 in /usr/local/lib/python3.7/dist-packages (from ftfy) (0.2.5)\n",
            "Thu May  5 00:59:04 2022       \n",
            "+-----------------------------------------------------------------------------+\n",
            "| NVIDIA-SMI 495.46       Driver Version: 460.32.03    CUDA Version: 11.2     |\n",
            "|-------------------------------+----------------------+----------------------+\n",
            "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
            "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
            "|                               |                      |               MIG M. |\n",
            "|===============================+======================+======================|\n",
            "|   0  Tesla V100-SXM2...  Off  | 00000000:00:04.0 Off |                    0 |\n",
            "| N/A   36C    P0    37W / 300W |  10927MiB / 16160MiB |      0%      Default |\n",
            "|                               |                      |                  N/A |\n",
            "+-------------------------------+----------------------+----------------------+\n",
            "                                                                               \n",
            "+-----------------------------------------------------------------------------+\n",
            "| Processes:                                                                  |\n",
            "|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |\n",
            "|        ID   ID                                                   Usage      |\n",
            "|=============================================================================|\n",
            "|  No running processes found                                                 |\n",
            "+-----------------------------------------------------------------------------+\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import json\n",
        "import os\n",
        "\n",
        "import numpy as np\n",
        "import openai\n",
        "import pandas as pd\n",
        "import pickle\n",
        "from sentence_transformers import SentenceTransformer\n",
        "from sentence_transformers import util as st_utils\n",
        "import torch\n",
        "\n",
        "openai.api_key = openai_api_key"
      ],
      "metadata": {
        "id": "pPaegbOraStd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# From: https://github.com/Deferf/CLIP_Video_Representation\n",
        "if not os.path.exists('MSRVTT_test_dict_CLIP_text.pt'):\n",
        "  !gdown 1-3tpfZzo1_D18WdrioQzc-iogEl-KSnA -O \"MSRVTT_test_dict_CLIP_text.pt\"\n",
        "if not os.path.exists('MSRVTT_test_dict_CLIP_visual.pt'):\n",
        "  !gdown 1Gp3_I_OvcKwjOQmn334-T4wfwQk29TCp -O \"MSRVTT_test_dict_CLIP_visual.pt\"\n",
        "if not os.path.exists('test_videodatainfo.json'):\n",
        "  !gdown 1BzTt1Bf-XJSUXxBfJVxLL3mYWLZ6odsw -O \"test_videodatainfo.json\"\n",
        "if not os.path.exists('JS_test_dict_CLIP_text.pt'):\n",
        "  !gdown --id 15mvFQxrWLNvBvFg4_9rr_Kqyzsy9dudj -O \"JS_test_dict_CLIP_text.pt\"\n",
        "\n",
        "# Load generated video transcriptions from Google cloud speed-to-text API.\n",
        "if not os.path.exists('video_id_to_gcloud_transcription_full.json'):\n",
        "  !gdown 1LTmvtf9zzw61O7D8YUqdS2mbql76nO6E -O \"video_id_to_gcloud_transcription_full.json\"\n",
        "\n",
        "# Load generated summaries from LM (comment this out to generate your own with GPT-3).\n",
        "if not os.path.exists('msr_full_summaries.pkl'):\n",
        "  !gdown 1ESXkRv3-3Kz1jZTNtkIhBXME6k1Jr9SW -O \"msr_full_summaries.pkl\""
      ],
      "metadata": {
        "id": "60WOfw5Fcv2-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Import helper functions from Portillo-Quintero et al. 2021\n",
        "!git clone https://github.com/Deferf/Experiments\n",
        "%cd Experiments\n",
        "from metrics import rank_at_k_precomputed,stack_encoded_dict,generate_sim_tensor,tensor_video_to_text_sim,tensor_text_to_video_metrics,normalize_matrix,pad_dict,list_recall\n",
        "%cd \"/content\""
      ],
      "metadata": {
        "id": "7v-nlqWbIEou",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "46c378cf-657c-42f3-d31b-2c13496eb102"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "fatal: destination path 'Experiments' already exists and is not an empty directory.\n",
            "/content/Experiments\n",
            "/content\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "##### Load RoBERTa (masked LM)"
      ],
      "metadata": {
        "id": "sdGpKAVDOuAg"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "roberta_model = SentenceTransformer('stsb-roberta-large').to(device)"
      ],
      "metadata": {
        "id": "8ORBat4VdysQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "##### Wrap GPT-3 (causal LM)"
      ],
      "metadata": {
        "id": "cGZfnNoTOyI1"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "gpt_version = \"text-davinci-002\"\n",
        "def prompt_llm(prompt, max_tokens=64, temperature=0, stop=None):\n",
        "  response = openai.Completion.create(engine=gpt_version, prompt=prompt, max_tokens=max_tokens, temperature=temperature, stop=stop)\n",
        "  return response[\"choices\"][0][\"text\"].strip()"
      ],
      "metadata": {
        "id": "HHoEoEP4zEfI"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## **Evaluate on MSR-Full**"
      ],
      "metadata": {
        "id": "uJ34kGQzwYL1"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Load raw text captions from MSR-Full.\n",
        "with open('test_videodatainfo.json', 'r') as j:\n",
        "  msr_full_info = json.loads(j.read())\n",
        "msr_full_vid_id_to_captions = {}\n",
        "for info in msr_full_info['sentences']:\n",
        "  if info['video_id'] not in msr_full_vid_id_to_captions:\n",
        "    msr_full_vid_id_to_captions[info['video_id']] = []\n",
        "  msr_full_vid_id_to_captions[info['video_id']].append(info['caption'])"
      ],
      "metadata": {
        "id": "ZzvX0Tw4vWLW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Reproduce original results with original eval code.\n",
        "msr_full_vid_id_to_clip_vid_feats = torch.load(\"/content/MSRVTT_test_dict_CLIP_visual.pt\", map_location=\"cpu\")\n",
        "msr_full_vid_ids_to_clip_text_feats = torch.load(\"/content/MSRVTT_test_dict_CLIP_text.pt\", map_location=\"cpu\")\n",
        "msr_full_vid_ids = list(msr_full_vid_ids_to_clip_text_feats.keys())\n",
        "msr_full_sim_tensor = generate_sim_tensor(msr_full_vid_ids_to_clip_text_feats, msr_full_vid_id_to_clip_vid_feats, msr_full_vid_ids)\n",
        "msr_full_vid_text_sim = tensor_video_to_text_sim(msr_full_sim_tensor)\n",
        "msr_full_metrics_vtt = rank_at_k_precomputed(msr_full_vid_text_sim)\n",
        "print(msr_full_metrics_vtt)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "M7GOzkn3uUTY",
        "outputId": "4368a5ff-fce4-4cbb-df9f-59c435bcfc4c"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "{'R@1': 40.301002502441406, 'R@5': 69.7324447631836, 'R@10': 79.19732666015625, 'Median_Rank': 2.0, 'Mean_Rank': 13.206688963210702, 'Std_Rank': 43.19973161311378}\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Transcription results from gCloud API.\n",
        "with open('video_id_to_gcloud_transcription_full.json', 'r') as j:\n",
        "  msr_full_vid_id_to_transcript = json.loads(j.read())\n",
        " \n",
        "# Sort video IDs by transcription length.\n",
        "num_transcripts = 0\n",
        "transcript_lengths = []\n",
        "for i in msr_full_vid_ids:\n",
        "  if msr_full_vid_id_to_transcript[i] is None:\n",
        "    transcript_lengths.append(0)\n",
        "  else:\n",
        "    num_transcripts += 1\n",
        "    transcript_lengths.append(len(msr_full_vid_id_to_transcript[i]))\n",
        "msr_full_sorted_vid_ids = [msr_full_vid_ids[i] for i in np.argsort(transcript_lengths)[::-1]]"
      ],
      "metadata": {
        "id": "SVWSgiY-xsm8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Summarize transcriptions with LLM.\n",
        "if os.path.exists('msr_full_summaries.pkl'):\n",
        "  msr_full_vid_id_to_summary = pickle.load(open('msr_full_summaries.pkl', 'rb'))\n",
        "else:\n",
        "\n",
        "  # Zero-shot LLM: summarize transcriptions.\n",
        "  msr_full_vid_id_to_summary = {}\n",
        "  for vid_id in msr_full_sorted_vid_ids:\n",
        "    transcript = msr_full_vid_id_to_transcript[vid_id]\n",
        "    print('Video ID:', vid_id)\n",
        "    print('Transcript:', transcript)\n",
        "  \n",
        "    if transcript is not None:\n",
        "      transcript = transcript.strip()\n",
        "      prompt = 'I am an intelligent video captioning bot.'\n",
        "      prompt += f'\\nI hear a person saying: \"{transcript}\".'\n",
        "      prompt += f\"\\nQ: What's a short video caption for this video? A: In this video,\"\n",
        "      print('Prompt:', prompt)\n",
        "      summary = prompt_llm(prompt, temperature=0, stop='.')\n",
        "      print('Summary:', summary)\n",
        "      msr_full_vid_id_to_summary[vid_id] = summary\n",
        "  \n",
        "    pickle.dump(msr_full_vid_id_to_summary, open(f'msr_full_summaries.pkl', 'wb'))"
      ],
      "metadata": {
        "id": "qe1w4FfbykGR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Compute RoBERTa features for all captions.\n",
        "msr_full_vid_id_to_roberta_feats = {}\n",
        "for vid_id in msr_full_sorted_vid_ids:\n",
        "  msr_full_vid_id_to_roberta_feats[vid_id] = roberta_model.encode(msr_full_vid_id_to_captions[vid_id], convert_to_tensor=True, device=device)"
      ],
      "metadata": {
        "id": "7ruc4aiNzTrk"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "topk = 100  # Pre-rank with top-100 from Portillo.\n",
        "combine_clip_roberta = True  # Combine CLIP (text-video) x RoBERTa (text-text) scores?\n",
        "portillo_vid_id_to_topk_vid_ids = {}\n",
        "socratic_vid_id_to_topk_vid_ids = {}\n",
        "msr_full_all_clip_text_feats = torch.cat([msr_full_vid_ids_to_clip_text_feats[i] for i in msr_full_sorted_vid_ids], dim=0).cpu().numpy()\n",
        "for vid_id in msr_full_sorted_vid_ids:\n",
        " \n",
        "  # Get Portillo top-K captions.\n",
        "  vid_feats = msr_full_vid_id_to_clip_vid_feats[vid_id]  # CLIP features for all frames of the video\n",
        "  vid_feat = normalize_matrix(torch.mean(vid_feats, dim = 0, keepdim = True)).cpu().numpy()\n",
        "  clip_scores = msr_full_all_clip_text_feats @ vid_feat.T\n",
        "  clip_scores = clip_scores.squeeze()\n",
        "  clip_scores = clip_scores.reshape(-1, 20)\n",
        "  clip_scores = np.max(clip_scores, axis=1)\n",
        "  sorted_idx = np.argsort(clip_scores).squeeze()[::-1]\n",
        "  portillo_topk_vid_ids = [msr_full_sorted_vid_ids[i] for i in sorted_idx[:topk]]\n",
        "  portillo_vid_id_to_topk_vid_ids[vid_id] = portillo_topk_vid_ids\n",
        "\n",
        "  # If no LLM summary, default to Portillo ranking.\n",
        "  socratic_vid_id_to_topk_vid_ids[vid_id] = portillo_topk_vid_ids\n",
        "  if vid_id not in msr_full_vid_id_to_summary:\n",
        "    continue\n",
        "\n",
        "  # Get RoBERTa scores between LLM summary and captions.\n",
        "  summary = msr_full_vid_id_to_summary[vid_id]\n",
        "  summary_feat = roberta_model.encode([summary], convert_to_tensor=True, device=device)\n",
        "  caption_feats = torch.cat([msr_full_vid_id_to_roberta_feats[i] for i in portillo_topk_vid_ids], dim=0)\n",
        "  roberta_scores = st_utils.pytorch_cos_sim(caption_feats, summary_feat).detach().cpu().numpy().squeeze()\n",
        "  roberta_scores = roberta_scores.reshape(-1, 20)\n",
        "  roberta_scores = np.max(roberta_scores, axis=1)\n",
        "\n",
        "  # Re-rank top-K with RoBERTa scores.\n",
        "  sort_idx = np.argsort(roberta_scores, kind='stable').squeeze()[::-1]\n",
        "  socratic_vid_id_to_topk_vid_ids[vid_id] = [portillo_topk_vid_ids[i] for i in sort_idx]\n",
        "\n",
        "  # Combine CLIP (text-video) x RoBERTa (text-text) scores.\n",
        "  if combine_clip_roberta:\n",
        "    clip_scores = np.sort(clip_scores, kind='stable').squeeze()[::-1][:topk]\n",
        "    scores = clip_scores * roberta_scores\n",
        "    sort_idx = np.argsort(scores, kind='stable').squeeze()[::-1]\n",
        "    socratic_vid_id_to_topk_vid_ids[vid_id] = [portillo_topk_vid_ids[i] for i in sort_idx]  # Override ranking from only LLM"
      ],
      "metadata": {
        "id": "y0LWrUoQzggH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Return R@1, R@5, R@10.\n",
        "def get_recall(vid_ids, socratic_subset, k=[1, 5, 10]):\n",
        " recall = []\n",
        " rank = []\n",
        " for vid_id in vid_ids:\n",
        "   sorted_vid_ids = portillo_vid_id_to_topk_vid_ids[vid_id]\n",
        "   if vid_id in socratic_subset:\n",
        "     sorted_vid_ids = socratic_vid_id_to_topk_vid_ids[vid_id]\n",
        "   recall.append([(vid_id in sorted_vid_ids[:i]) for i in k])\n",
        "   rank.append(sorted_vid_ids.index(vid_id) + 1 if vid_id in sorted_vid_ids else len(sorted_vid_ids))\n",
        " mdr = np.median(rank)\n",
        " return np.mean(np.float32(recall) * 100, axis=0), mdr\n",
        " \n",
        "subset_size = 1007  # Subset of long transcripts.\n",
        " \n",
        "# Portillo only.\n",
        "recall, mdr = get_recall(msr_full_sorted_vid_ids, msr_full_sorted_vid_ids[:0])\n",
        "print(f'R@1: {recall[0]:.1f}\\tR@5: {recall[1]:.1f}\\tR@10: {recall[2]:.1f}\\tMdR: {mdr}')\n",
        " \n",
        "# Socratic + Portillo.\n",
        "recall, mdr = get_recall(msr_full_sorted_vid_ids, msr_full_sorted_vid_ids[:subset_size])\n",
        "print(f'R@1: {recall[0]:.1f}\\tR@5: {recall[1]:.1f}\\tR@10: {recall[2]:.1f}\\tMdR: {mdr}')\n",
        " \n",
        "# Portillo only on long transcripts.\n",
        "recall, mdr = get_recall(msr_full_sorted_vid_ids[:subset_size], msr_full_sorted_vid_ids[:0])\n",
        "print(f'R@1: {recall[0]:.1f}\\tR@5: {recall[1]:.1f}\\tR@10: {recall[2]:.1f}\\tMdR: {mdr}')\n",
        " \n",
        "# Socratic + Portillo on long transcripts.\n",
        "recall, mdr = get_recall(msr_full_sorted_vid_ids[:subset_size], msr_full_sorted_vid_ids[:subset_size])\n",
        "print(f'R@1: {recall[0]:.1f}\\tR@5: {recall[1]:.1f}\\tR@10: {recall[2]:.1f}\\tMdR: {mdr}')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "n4sW3bvZ07z8",
        "outputId": "ea12c145-fdb9-44fd-a884-c11426b5c3e6"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "R@1: 40.2\tR@5: 69.7\tR@10: 79.2\tMdR: 2.0\n",
            "R@1: 44.7\tR@5: 71.2\tR@10: 80.0\tMdR: 2.0\n",
            "R@1: 41.5\tR@5: 69.6\tR@10: 77.4\tMdR: 2.0\n",
            "R@1: 54.9\tR@5: 74.0\tR@10: 79.9\tMdR: 1.0\n"
          ]
        }
      ]
    }
  ]
}